#
# Author: yizhigopher
# Description: 目前实现当日的调度问题，暂时不考虑联络线以及火电检修次数的问题
# File: electric_scheduling.py
# Date: 2023/11/16 15:13:06
#
import math

import numpy as np

from data.readData import TrainDataReader
from environment.storage import Storage

trainData = TrainDataReader()
alpha = 1
beta = 1
theta = 10
fire_rate = 1e-3
rate = 1


class PowerDayAheadSchedule:
    """
    :param require_ 传入负荷的时序数据 shape=[24, 1]
    :param fireP    传入火电出力限制 shape=[火电场站数, 4] 0->sudden upper bound 1-> normal upper bound 2-> sudden lower bound...
    :param fireExtra传入火电启停成本
    :param windP    传入风电预测出力, 要求我们的风电出力不能大于这个限制 shape=[24, 1]
    :param sunP     传入光电预测出力, 要求我们的光电出力不能大于这个限制 shape=[24, 1]
    :param waterP   传入水电预测出力, 要求我们的光电出力不能大于这个限制 shape=[水电场站数, 24]
    :param battery  化学电池储能信息
    :param saving   抽蓄储能信息
    定义每日的电力调度的环境
    其中，火力、需求、风电、光电的数据都是随时间变化而变化并制定限制
    尝试以(时间、火电出力，风电出力，水电出力, 光电出力，化学电池状态、水电储蓄状态)作为环境的状态信息
    """

    def __init__(self, d: int, require_: list, fireP: list, fireExtra: list, windP: list, sunP: list, waterP: list,
                 battery: Storage, saving: Storage, **kwargs):
        self.trainDuration = d
        self.fire_station_num = kwargs.get("fire-station-num")
        self.water_station_num = kwargs.get("water-station-num")

        # 记录火电启停成本
        self.fire_extra = fireExtra
        # 添加出力限制
        self.requirement_t = require_
        self.wind_limit_t = windP
        self.solar_limit_t = sunP
        self.fire_limit = fireP
        self.water_limit = waterP
        # 增加储能信息
        self.battery = battery
        self.saving = saving
        self.p_storage_t = 0

        self.curr_t = 0  # 当前时刻
        self.p_wind_t = 0  # 风电出力
        self.p_solr_t = 0  # 光电出力
        self.battery_t = 0  # 化学电池蓄能状态
        self.saving_t = 0  # 抽蓄储能出力
        self.p_water_t = [0 for _ in range(self.water_station_num)]  # 水电出力状态
        self.p_fire_t = [1 for _ in range(self.fire_station_num)]  # 上一时刻火电出力情况，用于计算启停成本
        self.p_total_t = 0

    def step(self, action):
        """
        :param action: 动作执行定义, 包括火电出力以及储能（化学电池和水电）
        :return: 状态信息, reward以及是否到达结束
        """
        # TODO：环境模型的reward该如何构建
        done = False
        # update state
        last_fire = self.p_fire_t
        self.p_fire_t = [x if x>0 else 0 for x in action["fire"]]
        self.p_water_t = [x if x>0 else 0 for x in action["water"]]
        self.p_wind_t = action['wind']
        self.p_solr_t = action['solar']
        # 不使用神经网络生成电池的出力情况，而是根据当前情况，采用贪心算法，生成储能状态
        # self.battery_t = action["cost-battery"]
        # self.saving_t = action["cost-saving"]

        self.p_total_t = 0
        for fire_out in self.p_fire_t:
            self.p_total_t += fire_out

        for water_out in self.p_water_t:
            self.p_total_t += water_out

        self.p_total_t += self.p_wind_t + self.p_solr_t

        cost_value = 0
        # 计算启停成本
        extra = 0
        for i in range(len(self.p_fire_t)):
            if (last_fire[i] == 0 and self.p_fire_t[i] > 0) or (last_fire[i] > 0 and self.p_fire_t[i] == 0):
                extra += self.fire_extra[i]
        cost_value -= extra
        if self.curr_t == self.trainDuration-1:
            done = True

        # 暂不考虑电池
        # self.handle_storage()

        if self.requirement_t[self.curr_t] > self.p_total_t:
            satisfiedPayloadReward = (self.p_total_t - self.requirement_t[self.curr_t])
        else:
            satisfiedPayloadReward = 0

        # if self.p_total_t > self.requirement_t[self.curr_t]:
        #     more = self.p_total_t - self.requirement_t[self.curr_t]
        #     self.battery_t += more / self.battery.eletric_rate  # 暂不考虑蓄电池能存多少电, 以及其使用寿命、

        # 如果风光电池出力策略超出预测范围，那么就扣除多余出力值
        energy_cost = 0
        if self.p_solr_t > self.solar_limit_t[self.curr_t]:
            energy_cost -= self.p_solr_t - self.solar_limit_t[self.curr_t]
        if self.p_wind_t > self.wind_limit_t[self.curr_t]:
            energy_cost -= self.p_wind_t - self.wind_limit_t[self.curr_t]

        # 计算水电出力超出约束部分的值
        for i in range(self.water_station_num):
            if self.p_water_t[i] > self.water_limit[i][self.curr_t]:
                cost_value -= (self.p_water_t[i] - self.water_limit[i][self.curr_t])
            if self.p_water_t[i] < 0:
                cost_value -= -self.p_water_t[i]

        # 计算超出火电出力范围的部分需要将超出的部分计入reward
        for i in range(len(self.p_fire_t)):
            # CHECK
            """
            分为两个部分，一个是瞬时额定功率，一个是连续额定功率
            如果发电功率能稳定在连续额定功率, 则不扣除reward
            如果在连续额定功率和瞬时额定功率之间, 则扣除0.1*功率
            如果超出瞬时额定功率, 则扣除想响应的功率
            """
            sudden_power = [self.fire_limit[i][2], self.fire_limit[i][0]]
            normal_power = [self.fire_limit[i][3], self.fire_limit[i][1]]

            if normal_power[1] > self.p_fire_t[i] > normal_power[0]:
                continue
            if self.p_fire_t[i] > sudden_power[1]:
                cost_value -= (self.p_fire_t[i] - sudden_power[1]) * 1
                # reward -= INF
                # done = True
                continue
            if self.p_fire_t[i] < sudden_power[0]:
                cost_value -= (sudden_power[0] - self.p_fire_t[i]) * 1
                # reward -= INF
                # done = True
                continue
            if self.p_fire_t[i] > normal_power[1]:
                cost_value -= (self.p_fire_t[i] - normal_power[1]) * 0.6
                continue
            if self.p_fire_t[i] < normal_power[0]:
                cost_value -= (normal_power[0] - self.p_fire_t[i]) * 0.6
                continue

        fire_cost = 0
        for i in range(len(self.p_fire_t)):
            p_fire_mw = self.p_fire_t[i]
            a, b, c = trainData.FireElectric[i][12:15]
            fire_cost += a * p_fire_mw ** 2 + b * p_fire_mw + c
            # fire_cost += self.p_fire_t[i]
        cost_value -= fire_cost * fire_rate  # 扣除火电发电成本

        reward = (alpha * satisfiedPayloadReward + cost_value * beta + energy_cost)*rate
        # reward += (self.p_solr_t + self.p_wind_t)

        self.curr_t = (self.curr_t + 1) % self.trainDuration
        return np.array([
            self.curr_t,
            self.requirement_t[self.curr_t],
            self.wind_limit_t[self.curr_t],
            self.solar_limit_t[self.curr_t],
            # self.battery_t,
            # self.saving_t
        ]), reward, done, {}

    def reset(self):
        self.curr_t = 0  # 当前时刻
        self.p_wind_t = self.wind_limit_t[self.curr_t]  # 风电出力
        self.p_solr_t = self.solar_limit_t[self.curr_t]  # 光电出力
        self.battery_t = 0  # 化学电池蓄能状态
        self.saving_t = 0
        self.p_storage_t = 0
        self.p_water_t = [0 for _ in range(len(self.water_limit))]  # 水电存储状态
        self.p_fire_t = [1 for _ in range(len(self.fire_limit))]
        self.p_total_t = 0
        return np.array([
            self.curr_t,
            self.requirement_t[self.curr_t],
            self.p_wind_t,
            self.p_solr_t,
            # self.battery_t,
            # self.saving_t
        ])

    def handle_storage(self):
        self.p_storage_t = 0
        if self.p_total_t < self.requirement_t[self.curr_t]:
            # 从化学储能中获取备份
            need = self.requirement_t[self.curr_t] - self.p_total_t
            if need / self.battery.eletric_rate <= self.battery_t:
                self.p_storage_t += need / self.battery.eletric_rate
                self.battery_t -= need / self.battery.eletric_rate
                self.p_total_t = self.requirement_t[self.curr_t]
            else:
                battery_output = self.battery_t * self.battery.eletric_rate
                need -= battery_output
                self.p_total_t += battery_output
                self.p_storage_t += self.battery_t
                self.battery_t = 0
                if need / self.saving.eletric_rate <= self.saving_t:
                    # 抽蓄储能电池足够满足需求
                    self.p_storage_t += need / self.saving.eletric_rate
                    self.saving_t -= need / self.saving.eletric_rate
                    self.p_total_t = self.requirement_t[self.curr_t]
                    self.p_total_t = self.requirement_t[self.curr_t]
                else:
                    saving_output = self.saving_t * self.saving.eletric_rate
                    need -= saving_output
                    self.p_total_t += saving_output
                    self.p_storage_t += self.saving_t
                    self.saving_t = 0
